from email.mime import audio
import os

import numpy as np
import pandas as pd
import librosa as lb
import matplotlib.pyplot as plt

from copy import copy
from glob import glob
from datasets.dataio import read_audio
from datasets.class_dict import classes_labels


class VisualSED(object):
    def __init__(
        self, 
        wav_path, 
        meta_path, 
        score_path, 
        out_path="./output/", 
        n_mels=128, 
        hop_length=256, 
        n_fft=2048,
        sr=16000,
        audio_len=10,
        decision_thd=0.5
    ):
        self.audio_path = out_path + "/audio/"
        self.fig_path = out_path + "/figs/"
        self.n_mels = n_mels
        self.hop_length = hop_length
        self.n_fft = n_fft
        self.sr = sr
        self.audio_len=audio_len
        self.thd = decision_thd
        
        self.wavs = glob(wav_path + "/*.wav")
        self.labels = list(pd.read_csv(meta_path, sep="\t").groupby("filename"))
        self.scores = glob(score_path + "/*.tsv")
        assert len(self.wavs) == len(self.labels), "Lengths of wav, label are not equal."
        self.idx2wav = {i: x.split("/")[-1].replace(".wav", "") for i, x in enumerate(self.wavs)}
        self.lab_idx = {x[0].replace(".wav", ""): i for i, x in enumerate(self.labels)}
        self.score_idx = {x.split("/")[-1].replace(".tsv", ""): i for i, x in enumerate(self.scores)}
        
        if not os.path.exists(self.fig_path):
            os.makedirs(self.fig_path)
        if not os.path.exists(self.audio_path):
            os.makedirs(self.audio_path)
    
    def __len__(self):
        return len(self.wavs)
    
    def _plot(self, wav_path, label, score_idx, wav_name):
        wav, _, _, _ = read_audio(wav_path, False, False, self.audio_len * self.sr)
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 8))
        plt.suptitle(f"{wav_name}")
        # Plot spectrogram
        spec = lb.feature.melspectrogram(y=wav.numpy(), sr=self.sr, n_fft=self.n_fft, n_mels=self.n_mels, hop_length=self.hop_length)
        ax1.imshow(lb.power_to_db(spec), aspect="auto", origin="lower", cmap="jet")
        ax1.set_title("LogMel spectrogram")
        ax1.set_ylabel("LogMel features")
        ax1.set_xticks(np.linspace(0, spec.shape[-1], 11), ["{:.1f}".format(x) for x in np.linspace(0, self.audio_len, 11)])
        
        # Plot labels
        labels_mat = np.zeros((len(classes_labels) * 2 + 1, spec.shape[-1]))

        if not pd.isna(label).any().any():
            for l in label.iterrows():
                l = l[1]
                lab_name, onset, offset = l["event_label"], l["onset"], l["offset"]
                lab_idx = classes_labels[lab_name]
                st = int(onset / self.audio_len * spec.shape[-1])
                ed = int(offset / self.audio_len * spec.shape[-1])
                labels_mat[lab_idx * 2 + 1, st:ed] = 1
            
        ax2.imshow(labels_mat, aspect="auto", origin="lower", cmap="gray")
        ax2.set_yticks(np.arange(1, len(classes_labels) * 2 + 1, 2), labels=list(classes_labels.keys()))
        ax2.set_xticks(np.linspace(0, spec.shape[-1], 11), ["{:.1f}".format(x) for x in np.linspace(0, self.audio_len, 11)])
        ax2.grid()
        
        
        score_mat = np.zeros((len(classes_labels) * 2 + 1, spec.shape[-1]))
        if score_idx is not None:
            score = pd.read_csv(self.scores[score_idx], sep="\t")
            for l in score.iterrows():
                l = l[1]
                for lab in (classes_labels.keys()):
                    lab_value, onset, offset = l[lab], l["onset"], l["offset"]
                    if lab_value > self.thd:
                        lab_idx = classes_labels[lab]
                        st = int(onset / self.audio_len * spec.shape[-1])
                        ed = int(offset / self.audio_len * spec.shape[-1])
                        score_mat[lab_idx * 2 + 1, st:ed] = 1
        ax3.imshow(score_mat, aspect="auto", origin="lower", cmap="gray")
        ax3.set_yticks(np.arange(1, len(classes_labels) * 2 + 1, 2), labels=list(classes_labels.keys()))
        ax3.grid()
        ax3.set_xticks(np.linspace(0, spec.shape[-1], 11), ["{:.1f}".format(x) for x in np.linspace(0, self.audio_len, 11)])
        ax3.set_xlabel("Time (s)")
        plt.savefig(f"{self.fig_path}/{wav_name}.png")
    
    def show(self, wav_name, listen=False):
        if isinstance(wav_name, int):
            assert wav_name < len(self), "Index out of range."
            wav_idx = copy(wav_name)
            wav_name = self.idx2wav[wav_name]
        else:
            wav_idx = list(self.idx2wav.values()).index(wav_name)
        
        wav_path = self.wavs[wav_idx]
        label = self.labels[self.lab_idx[wav_name]][1]
        try:
            score_idx = self.score_idx[wav_name]
        except:
            score_idx = None
        
        print("Wav name: ", wav_path, wav_name, "Score: ", score_idx)
        self._plot(wav_path, label, score_idx, wav_name)
        if listen:
            os.system(f"cp {wav_path} {self.audio_path}")
        
if __name__ == "__main__":
    wav_path = "/home/shaonian/Datasets/DCASE/dcase2021/dataset/audio/validation/validation_16k/"
    meta_path = "/home/shaonian/Datasets/DCASE/dcase2021/dataset/metadata/validation/validation.tsv"
    score_path = "/home/shaonian/SED/sssl_sed/codes/exp/da_exp_loss/MixDomain_SSL_Nodecay+MTWarmup50InDomain_300epoch/version_0/metrics_test/teacher/scenario1/scores/"
    
    visual_tool = VisualSED(wav_path, meta_path, score_path, "./visual_results/MixDomain_SSL_Nodecay+MTWarmup50InDomain_300epoch/teacher/")
    for i in range(100):
        visual_tool.show(i, listen=True)
    